package com.piece.core.elasticsearch.support;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.json.jackson.JacksonJsonpMapper;
import co.elastic.clients.transport.ElasticsearchTransport;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.piece.core.elasticsearch.properties.ElasticsearchProperties;
import com.piece.core.framework.util.string.StringUtil;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.PooledObjectFactory;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.http.HttpHost;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.CredentialsProvider;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
//import org.elasticsearch.client.RestHighLevelClient;
import org.springframework.util.Assert;
import java.util.ArrayList;
import java.util.List;

/**
 * elasticsearch连接池工厂
 */
public class EsClientPoolFactory implements PooledObjectFactory<ElasticsearchClient> {

    private ElasticsearchProperties elasticsearchProperties;

    public EsClientPoolFactory(ElasticsearchProperties elasticsearchProperties) {
        this.elasticsearchProperties = elasticsearchProperties;
    }

    @Override
    public void activateObject(PooledObject<ElasticsearchClient> pooledObject) throws Exception {

    }

    @Override
    public void destroyObject(PooledObject<ElasticsearchClient> pooledObject) throws Exception {
        pooledObject.getObject();
    }

    @Override
    public PooledObject<ElasticsearchClient> makeObject() throws Exception {
        List<HttpHost> httpHosts = new ArrayList<>();
        List<String> clusterNodes = elasticsearchProperties.getHosts();

        clusterNodes.forEach(node -> {
            try {
                String[] parts = StringUtil.split(node, ":");
                Assert.notNull(parts, "Must defined");
                Assert.state(parts.length == 2, "Must be defined as 'host:port'");
                httpHosts.add(new HttpHost(parts[0], Integer.parseInt(parts[1]), elasticsearchProperties.getSchema()));
            } catch (Exception e) {
                throw new IllegalStateException("Invalid ES nodes " + "property '" + node + "'", e);
            }
        });

        RestClientBuilder builder = RestClient.builder(httpHosts.toArray(new HttpHost[0]));
        ElasticsearchTransport transport = new RestClientTransport(getRestClient(builder, elasticsearchProperties), new JacksonJsonpMapper());
        ElasticsearchClient client = new ElasticsearchClient(transport);

        return new DefaultPooledObject<>(client);
    }

    private static RestClient getRestClient(RestClientBuilder builder, ElasticsearchProperties elasticsearchProperties) {
        // Callback used the default {@link RequestConfig} being set to the {@link CloseableHttpClient}
        builder.setRequestConfigCallback(requestConfigBuilder -> {
            requestConfigBuilder.setConnectTimeout(elasticsearchProperties.getConnectTimeout());
            requestConfigBuilder.setSocketTimeout(elasticsearchProperties.getSocketTimeout());
            requestConfigBuilder.setConnectionRequestTimeout(elasticsearchProperties.getConnectionRequestTimeout());
            return requestConfigBuilder;
        });

        // Callback used to customize the {@link CloseableHttpClient} instance used by a {@link RestClient} instance.
        builder.setHttpClientConfigCallback(httpClientBuilder -> {
            httpClientBuilder.setMaxConnTotal(elasticsearchProperties.getMaxConnectTotal());
            httpClientBuilder.setMaxConnPerRoute(elasticsearchProperties.getMaxConnectPerRoute());

            // Callback used the basic credential auth
            ElasticsearchProperties.Account account = elasticsearchProperties.getAccount();
            if (!StringUtil.isEmpty(account.getUsername()) && !StringUtil.isEmpty(account.getUsername())) {
                final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
                credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(account.getUsername(), account.getPassword()));
                httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
            }

            return httpClientBuilder;
        });

        return builder.build();
    }

//    @Override
//    public PooledObject<RestHighLevelClient> makeObject() throws Exception {
//        List<HttpHost> httpHosts = new ArrayList<>();
//        List<String> clusterNodes = elasticsearchProperties.getHosts();
//
//        clusterNodes.forEach(node -> {
//            try {
//                String[] parts = StringUtil.split(node, ":");
//                Assert.notNull(parts, "Must defined");
//                Assert.state(parts.length == 2, "Must be defined as 'host:port'");
//                httpHosts.add(new HttpHost(parts[0], Integer.parseInt(parts[1]), elasticsearchProperties.getSchema()));
//            } catch (Exception e) {
//                throw new IllegalStateException("Invalid ES nodes " + "property '" + node + "'", e);
//            }
//        });
//
//        RestClientBuilder builder = RestClient.builder(httpHosts.toArray(new HttpHost[0]));
//        RestHighLevelClient client = getRestHighLevelClient(builder, elasticsearchProperties);;
//
//        return new DefaultPooledObject<>(client);
//    }
//
//    private static RestHighLevelClient getRestHighLevelClient(RestClientBuilder builder, ElasticsearchProperties elasticsearchProperties) {
//        // Callback used the default {@link RequestConfig} being set to the {@link CloseableHttpClient}
//        builder.setRequestConfigCallback(requestConfigBuilder -> {
//            requestConfigBuilder.setConnectTimeout(elasticsearchProperties.getConnectTimeout());
//            requestConfigBuilder.setSocketTimeout(elasticsearchProperties.getSocketTimeout());
//            requestConfigBuilder.setConnectionRequestTimeout(elasticsearchProperties.getConnectionRequestTimeout());
//            return requestConfigBuilder;
//        });
//
//        // Callback used to customize the {@link CloseableHttpClient} instance used by a {@link RestClient} instance.
//        builder.setHttpClientConfigCallback(httpClientBuilder -> {
//            httpClientBuilder.setMaxConnTotal(elasticsearchProperties.getMaxConnectTotal());
//            httpClientBuilder.setMaxConnPerRoute(elasticsearchProperties.getMaxConnectPerRoute());
//
//            // Callback used the basic credential auth
//            ElasticsearchProperties.Account account = elasticsearchProperties.getAccount();
//            if (!StringUtil.isEmpty(account.getUsername()) && !StringUtil.isEmpty(account.getUsername())) {
//                final CredentialsProvider credentialsProvider = new BasicCredentialsProvider();
//
//                credentialsProvider.setCredentials(AuthScope.ANY, new UsernamePasswordCredentials(account.getUsername(), account.getPassword()));
//                httpClientBuilder.setDefaultCredentialsProvider(credentialsProvider);
//            }
//
//            return httpClientBuilder;
//        });
//
//        return new RestHighLevelClient(builder);
//    }

    @Override
    public void passivateObject(PooledObject<ElasticsearchClient> pooledObject) throws Exception {

    }

    @Override
    public boolean validateObject(PooledObject<ElasticsearchClient> pooledObject) {
        return true;
    }

    public ElasticsearchProperties getElasticsearchProperties() {
        return elasticsearchProperties;
    }
}
